package com.tools.common.security;

import com.tools.common.container.RowKit;
import com.tools.common.object.Note;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * CIDR 前缀树
 * */
@Note("CIDR 前缀树")
public final class CIDRPrefixTree {

    @Note("根节点")
    private final TrieNode root;

    @Note("读写锁保证并发性能 + 线程安全")
    private final ReadWriteLock lock;

    /* *******************************************************************************************
     *
     *          构造器
     *
     * ******************************************************************************************* */

    public CIDRPrefixTree() {
        this.root = new TrieNode();
        this.lock = new ReentrantReadWriteLock();
    }

    public CIDRPrefixTree(String cidr) {
        this();
        private_insert(cidr, this.root);
    }

    public CIDRPrefixTree(String... cidrArray) {
        this();
        private_insertBatch(cidrArray, this.root);
    }

    public CIDRPrefixTree(Iterable<String> cidrIterable) {
        this();
        private_insertBatch(cidrIterable, this.root);
    }

    /* *******************************************************************************************
     *
     *          核心方法
     *
     * ******************************************************************************************* */

    public boolean isEmpty() {
        this.lock.readLock().lock();
        try {
            return this.root.children.isEmpty();
        } finally {
            this.lock.readLock().unlock();
        }
    }

    @Note("插入 CIDR 表达式")
    public void insert(String cidr) {
        this.lock.writeLock().lock();
        try {
            private_insert(cidr, root);
        } finally {
            this.lock.writeLock().unlock();
        }
    }

    @Note("批量插入 CIDR 表达式")
    public void insertBatch(String... cidrArray) {
        this.lock.writeLock().lock();
        try {
            private_insertBatch(cidrArray, this.root);
        } finally {
            this.lock.writeLock().unlock();
        }
    }

    @Note("批量插入 CIDR 表达式")
    public void insertBatch(Iterable<String> cidrIterable) {
        this.lock.writeLock().lock();
        try {
            private_insertBatch(cidrIterable, this.root);
        } finally {
            this.lock.writeLock().unlock();
        }
    }


    @Note("删除 CIDR 表达式")
    public boolean remove(String cidr) {
        this.lock.writeLock().lock();
        try {
            return private_remove(cidr, this.root);
        } finally {
            this.lock.writeLock().unlock();
        }
    }

    public void removeAll(String... cidrArray) {
        this.lock.writeLock().lock();
        try {
            if(cidrArray != null && cidrArray.length > 0) {
                for(String cidr : cidrArray) {
                    private_remove(cidr, this.root);
                }
            }
        } finally {
            this.lock.writeLock().unlock();
        }
    }

    public void removeAll(Iterable<String> cidrIterable) {
        this.lock.writeLock().lock();
        try {
            if(cidrIterable != null) {
                for(String cidr : cidrIterable) {
                    private_remove(cidr, this.root);
                }
            }
        } finally {
            this.lock.writeLock().unlock();
        }
    }

    @Note("清空前缀树")
    public void clear() {
        this.lock.writeLock().lock();
        try {
            this.root.children = new ConcurrentHashMap<>();
            this.root.isEndOfPrefix = false;
        } finally {
            this.lock.writeLock().unlock();
        }
    }

    @Note("判断IP地址是否在 CIDR 范围内")
    public boolean contains(String ip) {
        this.lock.readLock().lock();
        try {
            return private_contains(ip, this.root);
        } finally {
            this.lock.readLock().unlock();
        }
    }

    @Note("判断数组里所有的 IP 地址是否都在 CIDR 范围内")
    public boolean containsAll(String... ipArray) {
        if(ipArray == null || ipArray.length == 0) return false;
        this.lock.readLock().lock();
        try {
            for (String ip : ipArray) {
                if(!private_contains(ip, this.root)) return false;
            }
            return true;
        } finally {
            this.lock.readLock().unlock();
        }
    }

    @Note("判断可迭代容器里所有的 IP 地址是否都在 CIDR 范围内")
    public boolean containsAll(Iterable<String> ipIterable) {
        int size = RowKit.getIterableSize(ipIterable);
        if(size == 0) return false;
        this.lock.readLock().lock();
        try {
            for (String ip : ipIterable) {
                if(!private_contains(ip, this.root)) return false;
            }
            return true;
        } finally {
            this.lock.readLock().unlock();
        }
    }


    @Note("判断数组里的是否存在任意一个 IP 地址在 CIDR 范围内")
    public boolean containsAny(String... ipArray) {
        if(ipArray == null || ipArray.length == 0) return false;
        this.lock.readLock().lock();
        try {
            for (String ip : ipArray) {
                if(private_contains(ip, this.root)) return true;
            }
            return false;
        } finally {
            this.lock.readLock().unlock();
        }
    }

    @Note("判断可迭代容器里的是否存在任意一个 IP 地址在 CIDR 范围内")
    public boolean containsAny(Iterable<String> ipIterable) {
        int size = RowKit.getIterableSize(ipIterable);
        if(size == 0) return false;
        this.lock.readLock().lock();
        try {
            for (String ip : ipIterable) {
                if(private_contains(ip, this.root)) return true;
            }
            return false;
        } finally {
            this.lock.readLock().unlock();
        }
    }

    /* *******************************************************************************************
     *
     *          私有逻辑
     *
     * ******************************************************************************************* */

    private static int private_checkAndGetPrefixLength(String[] parts) {
        if(parts.length >= 2) {
            return Integer.parseInt(parts[1]);
        }
        return 32;
    }

    @Note("插入操作")
    private static void private_insert(String cidr, TrieNode root) {
        if(cidr == null || cidr.isEmpty()) return;
        String[] parts = cidr.split("/");
        String ip = parts[0];
        int prefixLength = private_checkAndGetPrefixLength(parts);
        // 将IP地址转为32位二进制形式
        String binaryIp = SysKit.toIPBinaryString(ip);
        TrieNode node = root;
        for (int i = 0; i < prefixLength; i++) {
            char bit = binaryIp.charAt(i);
            node = node.children.computeIfAbsent(bit, k -> new TrieNode());
        }
        // 设置该节点为CIDR的结束位置
        node.isEndOfPrefix = true;
    }

    @Note("批量插入操作")
    private static void private_insertBatch(String[] cidrArray, TrieNode root) {
        if (cidrArray != null && cidrArray.length != 0) {
            for (String cidr : cidrArray) {
                private_insert(cidr, root);
            }
        }
    }

    @Note("批量插入操作")
    private static void private_insertBatch(Iterable<String> cidrIterable, TrieNode root) {
        if (cidrIterable != null) {
            for (String cidr : cidrIterable) {
                private_insert(cidr, root);
            }
        }
    }


    @Note("判断是否包含操作")
    private static boolean private_contains(String ip, TrieNode root) {
        String binaryIp = SysKit.toIPBinaryString(ip);
        TrieNode node = root;
        for (int i = 0; i < 32; i++) {
            char bit = binaryIp.charAt(i);
            node = node.children.get(bit);
            //如果某个节点不存在，说明该IP不在任何CIDR范围内
            if (node == null) {
                return false;
            }
            //如果找到完整的前缀匹配，说明该IP在CIDR范围内
            if (node.isEndOfPrefix) {
                return true;
            }
        }
        return false;
    }


    @Note("移除操作")
    private static boolean private_remove(String cidr, TrieNode root) {
        if(cidr == null || cidr.isEmpty()) return false;
        String[] parts = cidr.split("/");
        String ip = parts[0];
        int prefixLength = private_checkAndGetPrefixLength(parts);
        String binaryIp = SysKit.toIPBinaryString(ip);
        return private_removeRecursively(binaryIp, root, 0, prefixLength);
    }

    @Note("递归查找并移除")
    private static boolean private_removeRecursively(String binaryIp, TrieNode node, int depth, int prefixLength) {
        if(depth == prefixLength) {
            if (node.isEndOfPrefix) {
                node.isEndOfPrefix = false;
                return node.children.isEmpty();
            }
            return false;
        }
        char bit = binaryIp.charAt(depth);
        TrieNode child = node.children.get(bit);
        //CIDR 不存在
        if(child == null) {
            return false;
        }
        //递归下一层
        boolean shouldDeleteChild = private_removeRecursively(binaryIp, child, depth + 1, prefixLength);
        if(shouldDeleteChild) {
            node.children.remove(bit);
            //如果节点没有子节点并且不是前缀的末尾，就删除该节点
            return node.children.isEmpty() && !node.isEndOfPrefix;
        }
        return false;
    }

    /* *******************************************************************************************
     *
     *          静态内部类
     *
     * ******************************************************************************************* */

    @Note("树的前缀节点")
    private static class TrieNode {
        private Map<Character, TrieNode> children;
        private boolean isEndOfPrefix;

        private TrieNode() {
            this.children = new ConcurrentHashMap<>();
            this.isEndOfPrefix = false;
        }
    }
}

